"""IRGroupNode 合并优化 Pass

将连续的相同 tag 和 params 的 IRGroupNode 合并为一个节点，
以优化 IR 结构并提高代码生成效率。
"""

from typing import Dict, Any, Optional
from .base import BasePass, PassStage, PassContext
from ..ir import IRModule, IRGroupNode, IRNoneNode


class GroupMergePass(BasePass):
    """IRGroupNode 合并优化 Pass

    遍历 IR 模块中的节点，找到连续的相同的 IRGroupNode，
    将它们合并为一个节点以优化 IR 结构。
    """

    name: str = "group_merge"
    stage = PassStage.OPTIMIZATION

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__()
        self.config = config or {}

    def run(self, ir_module: IRModule, context: PassContext) -> IRModule:
        """运行 IRGroupNode 合并优化

        Args:
            ir_module: 输入的 IR 模块

        Returns:
            优化后的 IR 模块
        """
        if not ir_module or not ir_module.nodes:
            return ir_module

        self._merge_groups(ir_module)
        return ir_module

    def _merge_groups(self, ir_module: IRModule) -> None:
        """合并连续的相同 IRGroupNode

        Args:
            ir_module: IR 模块
        """
        if len(ir_module.nodes) < 2:
            return

        i = 0
        merged_count = 0
        while i < len(ir_module.nodes):
            current_node = ir_module.nodes[i]
            if not isinstance(current_node, IRGroupNode):
                i += 1
                continue

            # 合并连续的相同 IRGroupNode
            j = i + 1
            while j < len(ir_module.nodes):
                next_node = ir_module.nodes[j]
                if not isinstance(next_node, IRGroupNode):
                    j += 1
                    break

                # 检查是否为相同的 IRGroupNode
                if current_node != next_node:
                    break

                current_node.merge(next_node)
                ir_module.nodes[j] = IRNoneNode(
                    name=next_node.name,
                    source_line=next_node.source_line,
                )
                merged_count += 1
                j += 1

            i = j

        if merged_count > 0:
            # 过滤掉 IRNoneNode，并重新设置所有节点的 parent
            new_nodes = [
                node for node in ir_module.nodes if not isinstance(node, IRNoneNode)
            ]
            ir_module.nodes = new_nodes
            # 重新设置 parent 链，确保 get_ir_module() 能够正常工作
            for node in ir_module.nodes:
                node.set_parent(ir_module)
